# ============================================================
# Thomas König & Stefan Eschenwecker
# The European Court of Justice and legal European integration
# 
# This script produces Figure 3 in the main text.
# ============================================================

# load packages
library(rstan)
library(brms)
library(tidyverse)


# load data
CourtData <- read.csv("Court/Data/PreparedCourtData.csv") %>% 
  filter(all_not_applic == 0)

# load model
RegModel <- readRDS("Court/Models/UnequalPrecision.rds")

# define scenarios
supra_signals <- c("Strong PS", "Weak PS", "Uninfo", "Cont", "Weak ME", "Strong ME")
ms_signal_conditions <- c("-2", "obs", "+2")
ms_signal_intensities <- c("obs", "+2")

ScenarioData <- expand_grid(
  supra_signal = supra_signals,
  net_ms_pref_std = ms_signal_conditions,
  ms_signal_int_std = ms_signal_intensities
) %>% 
  mutate(supra_signal_int = case_when(
    supra_signal == "Uninfo" ~ "No Intensity",
    supra_signal %in% c("Weak PS", "Weak ME") ~ "Weak Intensity",
    supra_signal %in% c("Cont", "Strong PS", "Strong ME") ~ "Strong Intensity",
    .default = NA_character_
  )
)

# remove implausible scenarios
ScenarioData <- ScenarioData %>%
  filter(!(
    (net_ms_pref_std == "obs" & ms_signal_int_std != "obs") |
    (ms_signal_int_std == "obs" & net_ms_pref_std == "+2") |
    (ms_signal_int_std == "obs" & net_ms_pref_std == "-2")
  ))


# helper function to create scenarios within loop
generate_scenario <- function(supra, ms_cond, euint, msint, base_data) {
  scenario <- base_data
  scenario$supra_signal <- supra
  scenario$supra_signal_int <- euint
  
  if (ms_cond == "obs") {
    scenario$net_ms_pref_std <- base_data$net_ms_pref_std
    scenario$ms_signal_int_std <- base_data$ms_signal_int_std
  } else {
    shift_pos <- as.numeric(ms_cond)
    shift_int <- as.numeric(msint)
    scenario$net_ms_pref_std <- base_data$net_ms_pref_std + shift_pos
    scenario$ms_signal_int_std <- base_data$ms_signal_int_std + shift_int
  }
  
  return(scenario)
}

# list to store results
PredictedResults <- list()


# Predicted outcome probabilities incorporating case-specific uncertainty

# loop over scenarios and save predicted probabilities in list
# (array: ndraws x nobs x outcomes)
# (Hint: takes some time to run)
set.seed(1408)
for (i in seq_len(nrow(ScenarioData))) {
  supra <- ScenarioData$supra_signal[i]
  ms_cond <- ScenarioData$net_ms_pref_std[i]
  euint <- ScenarioData$supra_signal_int[i]
  msint <- ScenarioData$ms_signal_int_std[i]
  
  scenario_name <- paste(supra, ms_cond, sep = "_")
  scenario_data <- generate_scenario(supra, ms_cond, euint, msint, CourtData)
  
  PredictedResults[[scenario_name]] <- fitted(RegModel,
    re_formula = NULL,
    scale = "response",
    newdata = scenario_data,
    summary = F
  )
}

# collapse draws over nobs across observations
# (Hint: takes some time to run)
TidyResults <- lapply(PredictedResults, function(x) {
  data.frame(apply(x, c(2, 3), mean))
})

# generate plot

# open file to save output
pdf("Court/Plots/PredProbsUneqPrecision.pdf", height = 10, width = 8)
# set up frame and axes
par(mfrow = c(1, 1), mar = c(4, 1, 1, 1))
plot(1, 1,
  type = "n", xlim = c(-.2, 1), ylim = c(23, 0), axes = FALSE,
  xlab = "Predicted Probability", ylab = ""
)
axis(1, at = seq(0, 1, length.out = 5), labels = c("0%", "25%", "50%", "75%", "100%"))
ypos <- seq(1, 24, by = 1)[-seq(4, 24, by = 4)]
segments(0, ypos, 1, ypos, col = "grey")

# add labels including nobs
labels <- c(rep(c(
  "-2 SD", "Observed Net", "+2 SD"), 5)
)
text(-.2, ypos, labels, adj = 0)
main_labels <- c(
  "EU Both PS (N=581)", " EU Single PS (N=465)",
  "EU Ambivalent (N=1079)", "EU Contradictory (N=372)",
  "EU Single ME (N=555)", "EU Both ME (N=879)"
)
ypos_main <- seq(0, 23, by = 4)
text(-.2, ypos_main, main_labels, adj = 0, font = 2)
ypos_signal <- seq(2, 23, by = 4)
text(x = -.22, y = ypos_signal, "MS Signal", srt = 90, cex = .75)

# define colors and shapes for outcomes
colors <- c("darkred", "orange", "aquamarine3")
shapes <- c(18, 17, 16)

# loop over summary results
# (plot points + CIs + label)
for (i in seq_along(TidyResults)) {
  y <- ypos[i]
  num_outcomes <- ncol(TidyResults[[i]])
  vertical_shift <- seq(-0.2, 0.2, length.out = num_outcomes)

  for (j in seq_len(num_outcomes)) {
    points(mean(TidyResults[[i]][, j]), y + vertical_shift[j],
      pch = shapes[j], col = colors[j], cex = 1.75
    )
    segments(quantile(TidyResults[[i]][, j], 0.025), y + vertical_shift[j],
      quantile(TidyResults[[i]][, j], 0.975), y + vertical_shift[j],
      lwd = 2.5, col = colors[j]
    )
    text(mean(TidyResults[[i]][, j]), y + vertical_shift[j] - 0.5,
      round(mean(TidyResults[[i]][, j]) * 100),
      cex = 0.75, col = "black"
    )
  }
}
dev.off()
